import importlib
import sys
from typing import Dict

import ray

from lagent.schema import AgentMessage
from lagent.utils import load_class_from_string


class AsyncAgentRayActor:
    """异步 Ray Actor 封装类，用于分布式代理部署"""

    def __init__(
        self,
        config: Dict,  # 代理配置
        num_gpus: int,  # 分配的 GPU 数量
    ):
        """
        初始化异步 Ray Actor。

        Args:
            config: 代理配置字典
            num_gpus: 需要分配的 GPU 数量
        """
        # 从配置中获取代理类名
        cls_name = config.pop('type')
        python_path = config.pop('python_path', None)
        
        # 动态加载代理类
        cls_name = load_class_from_string(cls_name, python_path) if isinstance(
            cls_name, str) else cls_name
        
        # 创建 Ray Actor 类，并指定 GPU 资源
        AsyncAgentActor = ray.remote(num_gpus=num_gpus)(cls_name)
        
        # 实例化 Actor
        self.agent_actor = AsyncAgentActor.remote(**config)

    async def __call__(self, *message: AgentMessage, session_id=0, **kwargs):
        """
        异步调用代理处理消息。

        Args:
            message: 输入消息
            session_id: 会话ID
            kwargs: 其他参数

        Returns:
            代理响应
        """
        # 远程异步调用 Actor 的处理方法
        response = await self.agent_actor.__call__.remote(
            *message, session_id=session_id, **kwargs)
        return response


class AgentRayActor:
    """同步 Ray Actor 封装类，用于分布式代理部署"""

    def __init__(
        self,
        config: Dict,  # 代理配置
        num_gpus: int,  # 分配的 GPU 数量
    ):
        """
        初始化同步 Ray Actor。

        Args:
            config: 代理配置字典
            num_gpus: 需要分配的 GPU 数量
        """
        # 从配置中获取代理类名
        cls_name = config.pop('type')
        python_path = config.pop('python_path', None)
        
        # 动态加载代理类
        cls_name = load_class_from_string(cls_name, python_path) if isinstance(
            cls_name, str) else cls_name
        
        # 创建 Ray Actor 类，并指定 GPU 资源
        AgentActor = ray.remote(num_gpus=num_gpus)(cls_name)
        
        # 实例化 Actor
        self.agent_actor = AgentActor.remote(**config)

    def __call__(self, *message: AgentMessage, session_id=0, **kwargs):
        """
        同步调用代理处理消息。

        Args:
            message: 输入消息
            session_id: 会话ID
            kwargs: 其他参数

        Returns:
            代理响应
        """
        # 远程调用 Actor 的处理方法并等待结果
        response = self.agent_actor.__call__.remote(
            *message, session_id=session_id, **kwargs)
        return ray.get(response)  # 同步等待结果
